{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 04. Performance Profiling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can profile the performance with a 3D FDTD simulation:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from line_profiler import LineProfiler\n",
    "\n",
    "import fdtd\n",
    "import fdtd.backend as bd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Backend\n",
    "\n",
    "Let's profile the impact of the backend. These are the possible backends:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- ``numpy`` (defaults to float64 arrays)\n",
    "- ``torch`` (defaults to float64 tensors)\n",
    "- ``torch.float32``\n",
    "- ``torch.float64``\n",
    "- ``torch.cuda`` (defaults to float64 tensors)\n",
    "- ``torch.cuda.float32``\n",
    "- ``torch.cuda.float64``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdtd.set_backend(\"numpy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, the ``numpy`` backend is preferred for standard CPU calculations\n",
    "with `\"float64\"` precision as it is slightly faster than ``torch`` backend on CPU. However, a significant performance improvement can be obtained by choosing ``torch.cuda`` on large enough grids.\n",
    "\n",
    "Note that, in FDTD, ``float64`` precision is generally preferred over ``float32`` to ensure numerical stability and prevent numerical dispersion. If this is of no concern to you, you can opt for ``float32`` precision, which especially on a GPU might yield a significant performance boost."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WAVELENGTH = 1550e-9\n",
    "SPEED_LIGHT: float = 299_792_458.0  # [m/s] speed of light"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create FDTD Grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100\n",
    "\n",
    "grid = fdtd.Grid(\n",
    "    (N, N, N),\n",
    "    grid_spacing=0.05 * WAVELENGTH,\n",
    "    permittivity=1.0,\n",
    "    permeability=1.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "add boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x boundaries\n",
    "grid[0:10, :, :] = fdtd.PML(name=\"pml_xlow\")\n",
    "grid[-10:, :, :] = fdtd.PML(name=\"pml_xhigh\")\n",
    "\n",
    "# y boundaries\n",
    "grid[:, 0:10, :] = fdtd.PML(name=\"pml_ylow\")\n",
    "grid[:, -10:, :] = fdtd.PML(name=\"pml_yhigh\")\n",
    "\n",
    "# z boundaries\n",
    "grid[:, :, 0:10] = fdtd.PML(name=\"pml_zlow\")\n",
    "grid[:, :, -10:] = fdtd.PML(name=\"pml_zhigh\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "add sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid[10+N//10:10+N//10, :, :] = fdtd.PlaneSource(\n",
    "    period=WAVELENGTH / SPEED_LIGHT, name=\"source\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "add objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid[10+N//5:4*N//5-10, 10+N//5:4*N//5-10, 10+N//5:4*N//5-10] = fdtd.Object(permittivity=2.5, name=\"center_object\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "grid summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup LineProfiler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create and enable profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiler = LineProfiler()\n",
    "profiler.add_function(grid.update_E)\n",
    "profiler.enable()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "run simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid.run(50, progress_bar=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Profiler Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "print profiler summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiler.print_stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "grid.visualize(z=N//2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
